# -*- coding: utf-8 -*-
"""
Created on Mon Sep 23 13:46:27 2024

@author: 59567
"""
from ast import Interactive
import os
import random
import httpx  # 导入 httpx 库
import gradio as gr
import pickle
import pandas as pd

from pandas.core.methods.describe import describe_timestamp_1d

from plot_base import multi_grid_line, multi_macro_bar_or_line, item_timeline_seasonal, xy_match_same_freq_plot

from data_source import get_data_by_meta
from data_process import to_daily, to_monthly_mean, TjdSingleData, gen_xy_same_freq, cal_tztd_index

from utils import load_yaml_file
from typing import Dict
# now = datetime.now().strftime("%Y-%m-%d")

from volcenginesdkarkruntime import Ark
from kaleido.scopes.plotly import PlotlyScope

from docx import Document
from docx.shared import Inches

scope = PlotlyScope()

client = Ark(
    api_key="504bfd41-c285-4924-9107-74bb71846399",
    base_url="https://ark.cn-beijing.volces.com/api/v3",
    timeout=httpx.Timeout(timeout=1800),
)


def set_data_begin_time(df, meta):
    filter_date = max([meta[y_name][1] for y_name in df.columns])
    df.index = pd.to_datetime(df.index) # 转换为日期时间格式，方便后续处理
    filter_date = pd.to_datetime(filter_date) # 转换为日期时间格式，方便后续处理
    df = df[df.index > filter_date]
    return df


def pre_process_by_tab_name(names):
    chart_type = names[0]
    # 获取图表数据列名
    names = names[1:]
    # 根据列名获取数据
    df = df_macro.loc[:,names]
    # 删除全为NaN的列
    df = df.dropna(how='all')
    # 保留两位小数
    df = df.round(2)
    # 设置数据开始时间
    df = df.iloc[-4:,:]
    
    return df


def plot_item(tab_name, names, start_date, add_text):
    chart_type = names[0]
    # 获取图表数据列名
    names = names[1:]
    # 根据列名获取数据
    df = df_macro.loc[:,names]
    # 删除全为NaN的列
    df = df.dropna(how='all')
    # 保留两位小数
    df = df.round(2)
    # 设置数据开始时间
    df = set_data_begin_time(df, meta['meta'])
    start_date = pd.to_datetime(start_date) # 转换为日期时间格式，方便后续处理
    df = df[df.index > start_date]
    # 如果图表类型包含timeline
    if 'timeline' in chart_type:
        # 获取图表类型和频率
        chart_type, freq = chart_type.split('_')
        
        # 如果频率为D，则将数据转换为日数据
        if freq == 'D':
            year_dic = {y_name:to_daily(df[y_name]) for y_name in df.columns}
        # 如果频率为M，则将数据转换为月均数据
        elif freq == 'M':
            year_dic = {y_name:to_monthly_mean(df[y_name]) for y_name in df.columns}
        else:
            year_dic = None
            print('freq error')
        # 调用item_timeline_seasonal函数，生成图表
        b = item_timeline_seasonal(df, freq, year_dic)
    # 如果图表类型为bar或line
    elif chart_type in ['bar','line']:
        
        # 调用multi_macro_bar_or_line函数，生成图表
        b = multi_macro_bar_or_line(df, chart_type, tab_name, add_text)
    # 如果图表类型为multi_line
    elif chart_type == 'multi_line':
        # 调用multi_grid_line函数，生成图表
        b = multi_grid_line(df)
    return b


def insert_dataframe_to_docx(doc, df, tab):
    """
    将 DataFrame 数据插入到新创建的 Word 文档中。

    :param df: 要插入的 DataFrame 对象
    :param output_path: 生成的 Word 文档的保存路径
    """
    # 添加一个标题
    doc.add_heading(tab, 1)
    if '日期' not in df.columns:
        df.insert(0, '日期', df.index)
    # df.insert(0, '日期', df.index)
    # 创建一个表格，行数为 DataFrame 的行数加 1（包含表头），列数为 DataFrame 的列数
    table = doc.add_table(rows=df.shape[0] + 1, cols=df.shape[1])

    # 填充表头
    hdr_cells = table.rows[0].cells
    for col_idx, col_name in enumerate(df.columns):
        hdr_cells[col_idx].text = col_name

    # 填充表格数据
    for row_idx in range(df.shape[0]):
        row_cells = table.rows[row_idx + 1].cells
        for col_idx in range(df.shape[1]):
            row_cells[col_idx].text = str(df.iloc[row_idx, col_idx])
    return doc



def ana_by_llm(text_input):
    
    # 请求1
    completion = client.chat.completions.create(
        model="doubao-1-5-thinking-pro-250415",
        messages=[
            {"role": "system", "content": "你是一个经济分析师，你会根据给定的数据，给出准确的数据描述。不超过150字"},
            {"role": "user", "content": f"{text_input}"},
        ],
    )
    print(completion.choices[0].message.content)
    msg = completion.choices[0].message.content
    return msg


def ana_why_llm(text_input):
    # 请求1
    completion = client.chat.completions.create(
        model="doubao-1-5-thinking-pro-250415",
        messages=[
            {"role": "system", "content": "下面是一些经济数据描述，为什么数据是这样，未来如何展望。不超过300字"},
            {"role": "user", "content": f"{text_input}"},
        ],
    )
    print(completion.choices[0].message.content)
    msg = completion.choices[0].message.content
    return msg


# 读取配置信息，输出经济数据到plot文件夹里
meta: Dict = load_yaml_file('config.yaml')

use_cache = True
save_cache = False

if use_cache:
    df_macro = pd.read_csv('df_macro.csv', index_col=0)
else:
    df_macro = get_data_by_meta(meta['meta'])

if save_cache:
    df_macro.to_csv('df_macro.csv')
plot_ctrl = meta['plot_control']    
# plot_html(df_macro, meta)

class ShowData:
    def __init__(self):
        self.cols = df_macro.columns.to_list()
        self.plot_ctrl = plot_ctrl
        self.tabs = list(plot_ctrl.keys())[:9]
        self.tab_tab = {}
        self.plot_fig = {}
        self.text_desp = {}
        self.text_why = {}
        self.text_year = {}
        self.text_slider = {}
        self.choose_tab = {}
        self.if_add_text_tab = {}
        self.button_tab = {}
        self.toke_hint = {}
        self.img_path = {}
        self.cache_desp = {}
        self.cache_why = {}
        self.cache_df = {}
        self.fig_all = []
        
    def plot_all(self):
        self.fig_all = []
        start_date = '2012-01-01'
        for tab in self.tabs:
            b = plot_item(tab, self.plot_ctrl[tab], start_date, False)
            path = f'./plot/{tab}.png'
            b.write_image(path)
            self.img_path[tab] = path
            self.fig_all.append(b)
        return self.fig_all
    
    def ana_all_cache(self):
        if os.path.exists('LLM.pkl'):
            with open('LLM.pkl', 'rb') as f:
                req = pickle.load(f)
            print('缓存已加载')
            return req
        else:
            return self.ana_all()
    
    def ana_all(self):
        texts_why = []
        markdowns = []
        for tab in self.tabs:
            names = self.plot_ctrl[tab]
            df = pre_process_by_tab_name(names)
            self.cache_df[tab] = df
            df_text = df.to_string()
            text = ana_by_llm(df_text)
            markdowns.append(text)
            self.cache_desp[tab] = text
            why_text = ana_why_llm(text)
            self.cache_why[tab] = why_text
            texts_why.append(why_text)
        
        req = [*markdowns, *texts_why]
        # 保存缓存
        with open('LLM.pkl', 'wb') as f:
            pickle.dump(req, f)
        print('缓存已保存')
        
        return req
    
    def write_report(self, pth):
        cache_fig = {tab:self.fig_all[i] for i, tab in enumerate(self.tabs)}
        cache_desp = self.cache_desp
        cache_why = self.cache_why
        
        doc = Document()
        doc.add_heading('宏观经济数据分析', 0)
        for tab in self.tabs:
            df = self.cache_df[tab]
            # fig = cache_fig[tab] # 这里可以根据需要选择是否使用缓存的图片
            # fig2 = self.plot_fig[tab]
            # 添加表格
            doc = insert_dataframe_to_docx(doc, df, tab)
            # 添加描述
            desp = cache_desp[tab]
            doc.add_paragraph(desp)
            # 添加逻辑和展望
            why = cache_why[tab]
            # 添加图片
            # self.img_path[tab] = f"./plot/{tab}.png"
            # fig.write_image(self.img_path[tab])
            doc.add_picture(self.img_path[tab], width=Inches(6))
            doc.add_paragraph(why)
            print(tab)

        # 保存文档
        try:
            doc.save(pth)
            gr.Info("报告已保存到 {}".format(pth))
        except Exception as e:
            v = random.randint(0, 10000)
            random_pth = f'./plot/宏观经济形势报告{v}.docx'
            gr.Info("报告已保存到 {}".format(random_pth))
            print(e)
                
        print(8978)
    
    def ask_more(self, hint, desp):
        print('ask', hint, desp)
        completion = client.chat.completions.create(
        model="deepseek-r1-250120",
        messages=[
            {"role": "system", "content": "下面是一些经济数据描述，请按照经济逻辑分析下这些数据，未来如何展望。不超过300字"},
            {"role": "user", "content": f"{desp} {hint}"},
        ],
            )
        print(completion.choices[0].message.content)
        msg = completion.choices[0].message.content
        return msg
    
    def change_slider(self, i, tab, add_text):
        print('change', i)
        start_date = df_macro.index[i]
        fig = plot_item(tab, self.plot_ctrl[tab], start_date, add_text)
        fig.write_image(f'./plot/{tab}.png', scale=3)
        return gr.Text(value=start_date), fig
        
    
    def html(self):
        with gr.Blocks() as demo:
            gr.Markdown("# Auto Report with LLM")
            
            with gr.Row():
                bt_plot = gr.Button("1、数据")
                bt_ana = gr.Button("2、描述与分析")
                bt_ana_cache = gr.Button("2、读取缓存")
                bt_report = gr.Button("3、保存报告")
                pth = gr.Text(label='路径', value='./plot/宏观经济形势分析报告.docx')
                
            gr.Markdown("# 宏观经济数据分析")
            for tab in self.tabs:
                gr.Markdown(f"### {tab}")
                self.plot_fig[tab] = gr.Plot()
                
                with gr.Row():
                    self.text_year[tab] = gr.Text(label=f'{tab}起始日期', value=df_macro.index[-40], scale=1)
                    self.text_slider[tab] = gr.Slider(minimum=1, maximum=len(df_macro.index), step=1, value=len(df_macro.index)-40, label=f'(拖动改变{tab}起始日期)', scale=2)
                    self.tab_tab[tab] = gr.Text(label=f'{tab}',value=tab, scale=1)
                    self.if_add_text_tab[tab] = gr.Checkbox(label='是否添加文字标签', value=False, scale=1) # 这个功能还没写好，暂时不使
                
                with gr.Row():
                    self.toke_hint[tab] = gr.Text(label='提问', value=f"请帮我分析下{tab},文字控制在100字以内", scale=2)
                    self.button_tab[tab] = gr.Button("分析", scale=1)
                    self.choose_tab[tab] = gr.Checkbox(label='是否选用本次输出', value=True, scale=1)
                
                self.text_desp[tab] = gr.Markdown(label=f"{tab}分析结果", value="AI数据分析...")
                self.text_why[tab] = gr.Textbox(label=f"{tab}分析原因", value="AI经济研判...", interactive=True)
            
            desp_text = [text for _,text in self.text_desp.items()]
            why_text = [text for _,text in self.text_why.items()]
            llm_req = [*desp_text, *why_text]
            
            bt_plot.click(fn=self.plot_all,
                          outputs=[fig for _,fig in self.plot_fig.items()]
                          )
            bt_ana.click(fn=self.ana_all, 
                         outputs=llm_req
                         )
            bt_ana_cache.click(fn=self.ana_all_cache,
                               outputs=llm_req
                               )
            bt_report.click(fn=self.write_report, 
                            inputs=[pth]
                            )
            for tab in self.tabs:
                self.text_slider[tab].change(
                    fn=self.change_slider,
                    inputs=[self.text_slider[tab], self.tab_tab[tab], self.if_add_text_tab[tab]],
                    outputs=[self.text_year[tab], self.plot_fig[tab]]
                )
                self.button_tab[tab].click(
                    fn=self.ask_more, 
                    inputs=[self.toke_hint[tab], self.text_desp[tab]], 
                    outputs=[self.text_why[tab]],
                    )
            
        return demo
    
sd = ShowData()
sd.html().launch()